"""
merge_lora.py
-------------
把 LoRA adapter 合并到 base 模型，生成独立可用模型
每行都有注释，直接改路径即可
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# 基础模型路径
BASE_MODEL = r"D:\models\qwen\qwen\Qwen3-0___6B"

# LoRA adapter 路径
ADAPTER = r"./qwen-finetune-output/checkpoint-6"

# 合并后模型保存路径
OUT_DIR = r"./merged-model"

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)

# 加载 base 模型
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 加载 LoRA adapter
model = PeftModel.from_pretrained(model, ADAPTER)

# 合并 LoRA 到 base 模型
merged = model.merge_and_unload()

# 保存合并后的模型
merged.save_pretrained(OUT_DIR)
tokenizer.save_pretrained(OUT_DIR)

print("合并完成，模型已保存到：", OUT_DIR)
